# _*_ coding: utf-8 _*_
import os
import numpy as np
import matplotlib.pyplot as plt
from copy import deepcopy
from random import sample, randint, shuffle, uniform
from time import time, gmtime, strftime, localtime
import argparse

from VCMTSP import Instance

np.set_printoptions(suppress=True, linewidth=999, threshold=np.inf)
plt.rc('font', family='Times New Roman', weight='normal', size=12)

parser = argparse.ArgumentParser()
parser.add_argument("--case", required=True, type=str, help="The case name!")
parser.add_argument("-m", required=True, type=int, help="The number of salesmen!")
args = parser.parse_args()

m = args.m
case = args.case

instance = Instance(file=f'../../data/TSPLIB_VC/{case}.tsp/{case}.tsp.txt', m=m)
Distance = instance.distance  # distance matrix
Cost = instance.cost  # traveling cost of a single path
A = instance.A  # accessibility matrix
D = instance.D  # depot
V = deepcopy(instance.cities[1:]).tolist()  # city set, exclude depot


def make_dir(path):
    path = path.strip()
    path = path.rstrip("\\")
    isExists = os.path.exists(path)
    if not isExists:
        os.makedirs(path)
        return True
    else:
        print(path + ' already exists.')
        return False


def salesmen_of(city):
    return np.where(A[city, :] == 1)[0].tolist()


def cities_of(salesman):
    return np.where(A[:, salesman] == 1)[0].tolist()


def find_nearest(city, roll):
    tar, min_dis = -1, 1e99
    for c in roll:
        if min_dis > Distance[city][c]:
            min_dis = Distance[city][c]
            tar = c
    assert tar != -1
    return tar, min_dis


class Solution:
    def __init__(self, init='EMP'):
        self.routes = None
        self.length = None
        self.lengths = [0 for _ in range(m)]
        if init == 'EMP':
            return
            # initialize by optimal insertion
        elif init == 'OPT':
            self.routes = [[D, D] for _ in range(instance.m)]  # empty routes
            unvisited = deepcopy(V)
            while len(unvisited) > 0:
                city = sample(unvisited, 1)[0]
                unvisited.remove(city)
                self.optimal_insert(city)
        # initialize by greedy build
        elif init == 'GRD':
            self.routes = [[D] for _ in range(instance.m)]
            unvisited = deepcopy(V)
            while unvisited:
                for r in range(instance.m):
                    access_cities = cities_of(r)  # cities access to route r
                    candidates = set(unvisited) & set(access_cities)  # unvisited
                    if candidates:
                        nearest_city = find_nearest(self.routes[r][-1], roll=candidates)
                        self.routes[r].append(nearest_city)
                        unvisited.remove(nearest_city)
            for r in range(instance.m):
                self.routes[r].append(D)
        # initialize by random
        elif init == 'RND':
            self.routes = [[D] for _ in range(instance.m)]
            unvisited = deepcopy(V)
            while unvisited:
                city = sample(unvisited, 1)[0]
                unvisited.remove(city)
                salesmen = salesmen_of(city)
                r = sample(salesmen, 1)[0]
                self.routes[r].append(city)
            for r in range(instance.m):
                self.routes[r].append(D)
        else:
            print(f'Error. init={init}')
            quit(-1)
        self.calc_length()

    def calc_length(self):
        self.lengths = []
        for route in self.routes:
            self.lengths.append(Cost(route))
        self.length = sum(self.lengths)

    def optimal_insert(self, city):
        salesmen = salesmen_of(city)
        opt_r, opt_i, opt_cost = -1, -1, 1e99
        for r in salesmen:
            for i in range(1, len(self.routes[r])):
                p, q = self.routes[r][i - 1], self.routes[r][i]
                cost = Distance[p][city] + Distance[city][q] - Distance[p][q]
                if opt_cost > cost:
                    opt_r, opt_i, opt_cost = r, i, cost
        self.routes[opt_r].insert(opt_i, city)
        self.lengths[opt_r] += opt_cost
        self.length = sum(self.lengths)

    def two_opt(self):
        total_improve = 0
        for r, route in enumerate(self.routes):
            improve = 0
            for i in range(1, len(route) - 1):  # exclude depot
                _j, min_diff = -1, 1e99
                for j in range(i + 1, len(route) - 1):
                    diff = Distance[route[i]][route[j + 1]] + Distance[route[i - 1]][route[j]] \
                           - Distance[route[i - 1]][route[i]] - Distance[route[j]][route[j + 1]]
                    if min_diff > diff:
                        min_diff = diff
                        _j = j
                if min_diff < 0:
                    route = route[:i] + route[_j:i - 1:-1] + route[_j + 1:]
                    self.routes[r] = route
                    improve += min_diff
            if abs(improve) < 1e-10:
                improve = 0
            self.lengths[r] = self.lengths[r] + improve
            total_improve += improve
        self.length = sum(self.lengths)
        return total_improve


def make_neighbor(solution, p_c):
    neighbor_solution = Solution(init='EMP')
    neighbor_solution.routes = []
    remove = []  # removed cities
    for route in solution.routes:
        reserve_route = []
        for city in route:
            if uniform(0, 1) < p_c or city == D:
                reserve_route.append(city)
            else:
                remove.append(city)
        neighbor_solution.routes.append(reserve_route)
    shuffle(remove)
    for city in remove:
        neighbor_solution.optimal_insert(city)
    neighbor_solution.calc_length()
    return neighbor_solution


problem = 'VCMTSP'
method = 'ABC'
feature = 'base'
plt.figure(figsize=(10, 4))
max_epoch = 1
max_gen = 1000
NP = 100  # size of population
Pc = 0.8  # probability to deepcopy cities
limit_scout = 100
if __name__ == '__main__':
    date = strftime("%Y-%m-%d", localtime())
    save_path = f'./results/{date}/{case}/{m}'
    make_dir(save_path)
    save_name = f'{method}_{feature}'
    start_time = time()
    fitness = []
    optimum = Solution(init='RND')
    for epoch in range(1, max_epoch + 1):
        POP = []  # population
        # Initialization.
        for _ in range(NP):
            POP.append(Solution(init='RND'))
        # Progress
        pBest = deepcopy(min(POP, key=lambda _: _.length))
        gBest = deepcopy(min(POP, key=lambda _: _.length))
        gBest_list = []
        pBest_list = []
        pWorst_list = []
        Stag = [0 for _ in range(NP)]  # mark the stagnated iteration of each food source
        for gen in range(max_gen):
            # Employed bee phase.
            for i in range(NP):
                neighbor = make_neighbor(POP[i], Pc)
                if POP[i].length > neighbor.length:
                    POP[i] = deepcopy(neighbor)
                    Stag[i] = 0
                else:
                    Stag[i] = Stag[i] + 1
            # Onlooker bee phase.
            for i in range(NP):
                a, b = sample(list(range(0, NP)), 2)
                if POP[a].length > POP[b].length:
                    a, b = b, a
                neighbor = make_neighbor(POP[a], Pc)
                if POP[a].length > neighbor.length:
                    POP[a] = deepcopy(neighbor)
                    Stag[a] = 0
                else:
                    Stag[a] = Stag[a] + 1
            # Scout bee phase.
            for i in range(NP):
                if Stag[i] >= limit_scout:
                    POP[i] = deepcopy(Solution(init='OPT'))
                    Stag[i] = 0
            # Record gBest.
            pBest = deepcopy(min(POP, key=lambda _: _.length))
            if gBest.length > pBest.length:
                gBest = deepcopy(pBest)
            pBest_list.append(pBest.length)
            gBest_list.append(gBest.length)
            pWorst_list.append(max(POP, key=lambda _: _.length).length)
            print(f'problem: {problem} {case} m={m}, method: {method}, feature: {feature}, NP={NP}, Pc={Pc}. '
                  f'At epoch {epoch} gen {gen}. length: {round(gBest.length, 2)}', end='\r')

            plt.clf()
            plt.subplot(1, 2, 1)
            plt.plot(range(len(gBest_list)), gBest_list, label='gBest')
            plt.plot(range(len(pBest_list)), pBest_list, label='pBest')
            plt.plot(range(len(pWorst_list)), pWorst_list, label='pWorst')
            plt.legend()
            plt.title(f'total length: {gBest.length:.2f}')
            plt.subplot(1, 2, 2)
            instance.draw_2(plt, routes=gBest.routes)
            plt.tight_layout()
            plt.pause(0.001)
        plt.savefig(f'{save_path}/{save_name}_curve.png')
        plt.show()

        if optimum.length > gBest.length:
            optimum = deepcopy(gBest)
        fitness.append(gBest_list[-1])
        mode = 'w' if epoch == 1 else 'ab'
        # ---------> Change save file here! <---------
        with open(rf'{save_path}/{save_name}_fitness.txt', mode) as f:
            np.savetxt(f, np.array(gBest_list).reshape(1, len(gBest_list)))

    print(f'problem: {problem} {case} m={m}, method: {method}, feature: {feature}, NP={NP}, Pc={Pc}. '
          f'At epoch {max_epoch} gen {max_gen}. length: {round(np.array(fitness).mean(), 2)}')
    print(f"Time cost {strftime('%H:%M:%S', gmtime(time() - start_time))}")

